import numpy as np
import pandas as pd
import tensorflow as tf
import os
import numpy as np
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'

in_channels = 3
dataframe = pd.read_csv('./data/trn_per_500.csv', header=0)
X0 = dataframe.iloc[:, :2000]
indices = list(range(len(X0)))

X = X0.iloc[indices, :].astype(np.float32).reset_index().drop('index', axis=1).values
n_train = int(1 * len(X))
X_train = X[: n_train, :]
X_train -= np.mean(X_train, axis=1).reshape(-1, 1)
X_train /= (np.std(X_train, axis=1)+0.001).reshape(-1, 1)


def get_next_batch(batch_size, X):
    for i in range(0, len(X), batch_size):
        start = i
        end = min(len(X), i + batch_size)
        yield X[start:end, :]


model_path = './checkpoints_2000'
predictions = 'dense10/BiasAdd:0'
batch_size = 20
learning_rate = 0.001
ckpt = tf.train.get_checkpoint_state(model_path + '/')
saver = tf.train.import_meta_graph(ckpt.model_checkpoint_path + '.meta')

type_map = {'围网': 0, '拖网': 1, '刺网': 2}
type_map_rev = {v: k for k, v in type_map.items()}

P = []
with tf.Session() as sess:
    saver.restore(sess, ckpt.model_checkpoint_path)
    graph = tf.get_default_graph()

    img_placeholder = graph.get_tensor_by_name('input_image:0')
    labels = graph.get_tensor_by_name('labels:0')
    learning_rate_ph = graph.get_tensor_by_name('learning_rate:0')

    X_train_fold = X_train
    print('bef run')
    for bx in get_next_batch(batch_size, X_train_fold):
        predictions_val = sess.run([predictions],
                               feed_dict={img_placeholder: bx,
                                          learning_rate_ph: learning_rate})
        pred = np.argmax(predictions_val[0], axis=1)
        P += [pred]
        print(pred)

pred_all = np.array(P).reshape(-1)
sub = dataframe[['ship']]
sub['pred'] = pd.Series(pred_all)
sub['pred'] = sub['pred'].map(type_map_rev)
sub.to_csv('result0130_cnn_channel2.csv', index=None, header=None)
